//
//  tftp.c
//  TFTPDemo
//
//  Created by 陈冠杰 on 16/8/1.
//  Copyright © 2016年 EzioChen. All rights reserved.
//

#include <stdio.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <string.h>
#include <netdb.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <sys/select.h>

#include <netinet/in.h>
#include <arpa/inet.h>
#include <sys/time.h>

#include "tftp.h"





static int  isCancel = 0;
char        *lastTransferData;
int         lastTransferDatalength;



void tftp_interrupt(struct tftp_conn *tc){
    isCancel = 1;
    if (lastTransferData != NULL) {
        free(lastTransferData);
        lastTransferData = NULL;
    }
}


/**
 *  close the connection handle,i.3, delete out local state
 */
void tftp_close(struct tftp_conn *tc) {

    if (lastTransferData != NULL) {
        free(lastTransferData);
        lastTransferData = NULL;
    }
    
    if (!tc) {
        return;
    }
    
    if (tc->fp != NULL) {
        fclose(tc->fp);
    }
    close(tc->sock);
    free(tc->msgbuf);
    free(tc);
    tc = NULL;
}

struct tftp_conn *tftp_connect(int type,char *fname,char *mode,const char *hostname, const short port)
{
    struct addrinfo hints;
    struct addrinfo *res = NULL;
    struct tftp_conn *tc;
    
    if (!fname || !mode ||!hostname) {
        return NULL;
    }
    
    tc = malloc(sizeof(struct tftp_conn));
    
    if (type == TFTP_TYPE_PUT) {
        tc->msgbuf = malloc(BLOCK_SIZE_PUT+TFTP_DATA_HDR_LEN);
        memset(tc->msgbuf, 0, BLOCK_SIZE_PUT+TFTP_DATA_HDR_LEN);
    }
    else if (type == TFTP_TYPE_GET) {
        tc->msgbuf = malloc(BLOCK_SIZE_GET+TFTP_DATA_HDR_LEN);
        memset(tc->msgbuf, 0, BLOCK_SIZE_GET+TFTP_DATA_HDR_LEN);
    }
    else{
        tc->msgbuf = malloc(BLOCK_SIZE_GET+TFTP_DATA_HDR_LEN);
        memset(tc->msgbuf, 0, BLOCK_SIZE_GET+TFTP_DATA_HDR_LEN);
    }
    
    
    
    if (tc == NULL) {
        return NULL;
    }
    
    isCancel = 0;
    
    /**
     Create a socket.
     check return value.
     */
    /*Specify communication domain (protocol family) and communication semantics (socket type )*/
    
    if ((tc->sock = socket(AF_INET, SOCK_DGRAM, 0)) < 0) {
        fprintf(stderr, "Unable to create Socket...\n");
        free(tc->msgbuf);
        free(tc);
        return NULL;
    }
    
    /* check the type = read or write file */
    if (type == TFTP_TYPE_PUT) {
        tc->fp = fopen(fname, "rb");
    }else if(type == TFTP_TYPE_GET){
        tc->fp = fopen(fname, "wb");
    }else{
        fprintf(stderr, "Invalid TFTP mode,must be put or get \n");
        return NULL;
    }
    /*  initial settings for get addrinfo function  */
    memset(&hints, 0, sizeof(hints));
    hints.ai_family = PF_UNSPEC;
    hints.ai_socktype = SOCK_DGRAM;
    char port_str[8];
    sprintf(port_str, "%d", port);
    
    /*
     * get address from host name
     * if error, gracefully clear up
     */
    
    /* host address information is stored in the varible res */
    int errocCode = getaddrinfo(hostname, port_str, &hints, &res);
    if (errocCode != 0) {
        fprintf(stderr, "Unable to get host address...\n");
        close(tc->sock);
        free(tc->msgbuf);
        free(tc);
        return NULL;
    }
    
    /* Assign address to the connection handle
     * You can assume that the first address in the hostent 
     * struct is the corrent one
     */
    /*initial settings for varsible tc which will be used in the later transmition */
    memcpy(&tc->peer_addr, res->ai_addr, res->ai_addrlen);
    tc->addrlen = sizeof(struct sockaddr_in);
    tc->type = type;
    tc->mode = mode;
    tc->fname = fname;
    tc->blocknr = 0;
    
    return tc;
}


/*
 Send a read request to the Server.
 1. Format message.
 2. Send the request using the connection handle.
 3. Retrun the number of bytes sent, or negative on error.
 */

int tftp_send_rrq(struct tftp_conn *tc) {
    char *blockChar = "blksize";
    char *blocksize =  "16384";
    
    /* struct tftp_rrq *rrq; */
    struct tftp_rrq* pTftp_rrq = malloc(TFTP_RRQ_LEN((*tc).fname,  (*tc).mode, blockChar, blocksize));
   
    memset(pTftp_rrq, 0, TFTP_RRQ_LEN((*tc).fname,(*tc).mode,blockChar,blocksize));
    
    pTftp_rrq->opcode = htons(OPCODE_RRQ);//transform host ascii to net ascii
    memcpy(pTftp_rrq->req, tc->fname, strlen(tc->fname));
    memcpy(pTftp_rrq->req + strlen(tc->fname) + 1, tc->mode,strlen(tc->mode));
    memcpy(pTftp_rrq->req + strlen(tc->fname) + 1 +strlen(tc->mode) +1 , blockChar, strlen(blockChar));
    memcpy(pTftp_rrq->req + strlen(tc->fname) + 1 +strlen(tc->mode) +1 + strlen(blockChar) +1, blocksize, strlen(blocksize));
    
    
    printf("tftp_rrq: %s\n",pTftp_rrq->req);
    printf("tftp_rrq: %s\n",tc->fname);
    printf("tftp_rrq: %s\n",tc->mode);
    printf("tftp_rrq: %s\n",blocksize);
    
    /* send read request to the server */
    int rqtResult = (int)sendto(tc->sock, pTftp_rrq,
                           TFTP_RRQ_LEN((*tc).fname,(*tc).mode,blockChar,blocksize), 0,
                           (struct sockaddr*) (&(tc->peer_addr)), (*tc).addrlen);
    if (rqtResult < 0)
        perror("sendto");
    
    free(pTftp_rrq);
    return rqtResult;  
}

/*
 * Send a write request to the server.
 1. Format message.
 2. Send the request using the connection handle.
 3. Retrun the number of bytes sent,or negative on error.
 */
int tftp_send_wrq(struct tftp_conn *tc) {
    
    char *blockChar = "blksize";
    char *blocksize =  "512";
    
    /* struct tftp_wrq *wrq; */
    struct tftp_wrq* pTftp_wrq = malloc(TFTP_WRQ_LEN(tc->fname, tc->mode,blockChar,blocksize));
    
    memset(pTftp_wrq, 0, TFTP_WRQ_LEN(tc->fname, tc->mode,blockChar,blocksize));
    pTftp_wrq->opcode = htons(OPCODE_WRQ);//transform host ascii to net ascii
    
    memcpy(pTftp_wrq->req, tc->fname, strlen(tc->fname));
    memcpy(pTftp_wrq->req + strlen(tc->fname) + 1, tc->mode, strlen(tc->mode));
    memcpy(pTftp_wrq->req + strlen(tc->fname) + 1 +strlen(tc->mode) +1 , blockChar, strlen(blockChar));
    memcpy(pTftp_wrq->req + strlen(tc->fname) + 1 +strlen(tc->mode) +1 + strlen(blockChar) +1, blocksize, strlen(blocksize));
    
    /* send write request to the server */
    int rqtResult = (int)sendto(tc->sock, pTftp_wrq, TFTP_WRQ_LEN(tc->fname, tc->mode,blockChar,blocksize), 0,
                           (struct sockaddr*) (&(tc->peer_addr)), (*tc).addrlen);
    
    if (rqtResult < 0)
        perror("sendto");
    
    free(pTftp_wrq);
    return rqtResult;
    
}



/*
 Acknowledge reception of a block.
 1. Format message.
 2. Send the acknowledgement using the connection handle.
 3. Return the number of bytes sent, or negative on error.
 */
int tftp_send_ack(struct tftp_conn *tc) {
    /* struct tftp_ack *ack; */
    struct tftp_ack* pTftp_ack = (struct tftp_ack*) malloc(TFTP_ACK_HDR_LEN);
    pTftp_ack->opcode = htons(OPCODE_ACK);//transform host ascii to net ascii
    /* Get receiving data block number from tc->msgbuf and ack to the server */
    tc->blocknr = ntohs(((struct tftp_data*) tc->msgbuf)->blocknr);
    pTftp_ack->blocknr = htons(tc->blocknr);
    
    int result = (int)sendto(tc->sock, pTftp_ack,TFTP_ACK_HDR_LEN, 0, (struct sockaddr*)(&(tc->peer_addr)), (*tc).addrlen);
    
    free(pTftp_ack);
    return result;
}


/*
 * Send a data block to the other side.
 *1. Format message.
 *2. Add data block to message according to length argument.
 *3. Send the data block message using the connection handle.
 *4. Return the number of bytes sent, or negative on error.
 
  You need to be able to resend data in case of a timeout. When resending, 
  the old message should be sent again and therefore on new message should be created.
 This can ,for example ,be handled by passing a negative length indicating that the creation of a new message
 should be skipped.
 */
int tftp_send_data(struct tftp_conn* tc,int length){
    /* struct tftp_data *tdata */
    struct tftp_data *pTftp_data = malloc(TFTP_DATA_HDR_LEN + BLOCK_SIZE_PUT);
    memset(pTftp_data, 0, TFTP_DATA_HDR_LEN + length);
    pTftp_data->opcode = htons(OPCODE_DATA);
    
    /* Get receiving ack number from tc->msgbuf and send data to the server */
    /* Retransmitold block in case of packet loss */
    if(tc->blocknr==1+ntohs(((struct tftp_ack*)tc->msgbuf)->blocknr))
    {
        pTftp_data->blocknr=htons(tc->blocknr);
        memcpy(pTftp_data->data,lastTransferData,lastTransferDatalength);
        
        if (lastTransferData != NULL) {
            free(lastTransferData);
            lastTransferData = NULL;
        }
        
        /* transmit last transfered data */
        //printf("1==>%d\n",tc->blocknr);
        int result=(int)sendto(tc->sock, pTftp_data, TFTP_DATA_HDR_LEN+lastTransferDatalength,
                          0,(struct sockaddr*)(&(tc->peer_addr)), tc->addrlen);
        
        free(pTftp_data);
        return result;
    }
    /* Transmit new block */
    else{
        tc->blocknr=ntohs(((struct tftp_ack*)tc->msgbuf)->blocknr)+1;
        pTftp_data->blocknr= htons(tc->blocknr);
        
        /* read data from local file system and send it to the server */
        if(!feof(tc->fw))
        {
            int dataLength = (int)fread(pTftp_data->data,1,length,tc->fw);
            //printf("2==>%d\n",pTftp_data->blocknr);
            int result=(int)sendto(tc->sock, pTftp_data, TFTP_DATA_HDR_LEN+dataLength, 0,(struct sockaddr*)(&(tc->peer_addr)), tc->addrlen);
            
            /* record data as lastTranserData in case of packet loss */
            if (lastTransferData != NULL) {
                free(lastTransferData);
                lastTransferData = NULL;
            }
            
            lastTransferData=malloc(dataLength);
            memset(lastTransferData,0,dataLength);
            memcpy(lastTransferData,pTftp_data->data,dataLength);
            lastTransferDatalength=dataLength;
            
            free(pTftp_data);
            return result;
        }
        
        if (lastTransferData != NULL) {
            free(lastTransferData);
            lastTransferData = NULL;
        }
        free(pTftp_data);
    }
    return -1;
}


/*
 Transfer a file to or from the server.
 */
int tftp_transfer(struct tftp_conn *tc) {
    
    int retval = 0;         //Transfer result
    int len;                //Transfer data length
    int totlen = 0;         //Transfer total length
    int recvResult = 0;     //Length of reveiving data
    int sendResult = 0;     //Length of sending data
    struct timeval timeout; //Time out setting
    fd_set fdst;            //File descriptor for select function
    u_int16_t errorCode;    //Error code
    int isEnd=0;            //Transfer is end or not: 0-not, 1-end
    
    /* Sanity check */
    if (!tc) return -1;
    
    len = 0;
    
    /* After the connection request we should start receiving data
     * immediately */
    /* Check if we are putting a file or getting a file and send
     * the corresponding request. */
    if (TFTP_TYPE_GET == tc->type) {
        tftp_send_rrq(tc);
    } else if (TFTP_TYPE_PUT == tc->type) {
        tftp_send_wrq(tc);
    } else {
        return -1;
    }
    
    
    /*
     Put or get the file, block by block, in a loop.
     */
    do {
        /* 1. Wait for something from the server (using
         * 'select'). If a timeout occurs, resend last block
         * or ack depending on whether we are in put or get
         * mode. */
        while (1) {
            
            /*--- 中断 ---*/
            if (isCancel) {
                retval = -1;
                break;
            }
            
            
            /* Set a timeout for resending data. */
            timeout.tv_sec = TFTP_TIMEOUT;
            timeout.tv_usec = 0;
            
            FD_ZERO(&fdst);//Clear file descriptor set to enable status checking
            FD_SET(tc->sock, &fdst);//add fiel descriptor to the set
            /* Check status of values in the fdst to see if any fd is ready for reading or writing */
            int sltRlt = select(tc->sock + 1, &fdst, NULL, NULL, &timeout);
            if (sltRlt < 0) { //Error in the select function
                return -1;
            } else if (sltRlt == 0) { //Time out in the selection function, no readable, writable or error fds
                //******Retransmit********//
                if (TFTP_TYPE_GET == tc->type)
                {
                    /* Retransmit read request */
                    if(0==tc->blocknr)
                    {
                       tftp_send_rrq(tc);
                    }
                    /* Retransmit ack packet */
                    else{
                        printf("%s,%d\n", "Retransmit Ack", tc->blocknr);
                        tftp_send_ack(tc);
                    }
                }
                else if (TFTP_TYPE_PUT == tc->type)
                {
                    /* Retransmit write request */
                    if(0==tc->blocknr)
                    {
                        tftp_send_wrq(tc);
                    }
                    /* Retransmit data packet */
                    else{
                        printf("%s,%d\n", "Retransmit Data", tc->blocknr);
                        sendResult=tftp_send_data(tc,BLOCK_SIZE_PUT);
                    }
                } else {
                    return -1;
                }
                continue;//Continue to wait for receiving data
            }
            else {//readable, writable or error fds exist
                /* receive data  */
                recvResult = recvfrom(tc->sock,
                                      tc->msgbuf,       /**** 重点！重点！重点！****/
                                      1024*100, 0,      /**** <<------------ ****/
                                      (struct sockaddr*)(&(tc->peer_addr)),
                                      &tc->addrlen);
                break;
            }
        }
        
        /*--- 中断 ---*/
        if (isCancel) {
            retval = -1;
            goto out;
        }
        
        u_int16_t msgType = ntohs(((struct tftp_msg*)(tc->msgbuf))->opcode);//Msg type
        int serverSendBlkNr=0; //receiving data block number sent by server
        int serverAckBlkNr=0;  //receiving ack block number sent by server
        
        /* 2. Check the message type and take the necessary action. */
        switch (msgType) {
                
            case OPCODE_DATA:{
                /* Received data block, send ack */
                serverSendBlkNr=ntohs(((struct tftp_data*) tc->msgbuf)->blocknr);
                /* receive correct data packet and discard duplicated dat packets */
                if(tc->blocknr == serverSendBlkNr-1)
                {
                    /* write correct data to file */
                    int wd = (int)fwrite(tc->msgbuf+ TFTP_DATA_HDR_LEN, 1, recvResult-TFTP_DATA_HDR_LEN, tc->fw);
                    
                    totlen += recvResult - TFTP_DATA_HDR_LEN;
                    printf("Get => %d,%lu,%d writed:%d\n",serverSendBlkNr,recvResult - TFTP_DATA_HDR_LEN,totlen,wd);
                }
                /* Send ack packet to server */
                tftp_send_ack(tc);
                /* if length of receiving data packet is less than 516(512+4) bytes,
                 then the connection should be terminated */
                if(recvResult<BLOCK_SIZE_GET + 4 )
                {
                    isEnd=1;
                    printf("End of receiving data.\n");
                }
            }break;
                
            case OPCODE_ACK:{

                /* Received ACK, send next block */
                serverAckBlkNr = ntohs(((struct tftp_ack*)tc->msgbuf)->blocknr);
                
                /* receive correct ack packet and discard duplicated ack packets */
                if(tc->blocknr == serverAckBlkNr && serverAckBlkNr>0)
                {
                    totlen+=sendResult-TFTP_DATA_HDR_LEN;
                    
                    printf("Put => %d,%lu,%d\n",serverAckBlkNr,sendResult-TFTP_DATA_HDR_LEN,totlen);
                    /* if length of sending data packet is less than 516(512+4) bytes,
                     then the connection should be terminated */
                    if(sendResult< BLOCK_SIZE_PUT + 4 && serverAckBlkNr>0)
                    {
                        isEnd=1;
                        printf("End of sending data.\n");
                        break;
                    }
                }
                /* Send data packet to server */
                //u_int16_t blr = htons(ntohs(((struct tftp_ack*)tc->msgbuf)->blocknr)-0x626c);
                //memcpy(tc->msgbuf+2, &blr, 2);
                sendResult=tftp_send_data(tc,BLOCK_SIZE_PUT);
            }break;
                
            case OPCODE_ERR:{
                
                /* Handle error... */
                errorCode=ntohs(((struct tftp_err*) (tc->msgbuf))->errcode);
                fprintf(stderr,"%s", err_codes[errorCode]);
                retval=-1;
                goto out;  
            }break;
                
            case OPCODE_OACK:{
                /* Send ack packet to server */
                struct tftp_ack* pTftp_ack = (struct tftp_ack*)malloc(TFTP_ACK_HDR_LEN);
                pTftp_ack->opcode = htons(OPCODE_ACK);//transform host ascii to net ascii
                pTftp_ack->blocknr = htons(tc->blocknr);
                sendto(tc->sock, pTftp_ack,TFTP_ACK_HDR_LEN, 0, (struct sockaddr*)(&(tc->peer_addr)), (*tc).addrlen);
                free(pTftp_ack);
                
                if (tc->type == TFTP_TYPE_PUT) {
                    u_int16_t blr = htons(ntohs(((struct tftp_ack*)tc->msgbuf)->blocknr)-0x626c);
                    memcpy(tc->msgbuf+2, &blr, 2);
                    sendResult=tftp_send_data(tc,BLOCK_SIZE_PUT);
                }

            }break;
                
            default:{
                fprintf(stderr, "\nUnknown message type\n");  
                goto out;
            }
        }  
        
    } while (isEnd==0);  
    
    printf("\nTotal data bytes sent/received: %d.\n", totlen);
    
out:
    fclose(tc->fw);
    return retval;
    
}  


